from dynamiq.connections import E2B
from dynamiq.nodes.agents import Agent
from dynamiq.nodes.tools.e2b_sandbox import E2BInterpreterTool
from dynamiq.nodes.types import InferenceMode
from dynamiq.utils.logger import logger
from examples.llm_setup import setup_llm

AGENT_ROLE = """
Senior Data Scientist and Programmer with the ability to write well-structured Python code.
You have access to Python tools and the web to search for the best solutions to problems.
Generally, you follow these rules:
    - ALWAYS FORMAT YOUR RESPONSE IN MARKDOWN.
    - Use double quotes for property names.
    - Ensure the code is correct and runnable; reiterate if it does not work.
{{ my_role }}
{% raw %}
{% endraw %}
"""


PROMPT = """
Calc x^2 for x = 2
"""


def create_agent():
    """
    Create and configure the agent with necessary tools.

    Returns:
        Agent: A configured Dynamiq Agent ready to run.
    """
    tool = E2BInterpreterTool(connection=E2B())
    llm = setup_llm(model_provider="gpt", model_name="gpt-4o-mini", temperature=0.001)

    agent = Agent(
        name="React Agent",
        llm=llm,
        tools=[tool],
        role=AGENT_ROLE,
        max_loops=10,
        inference_mode=InferenceMode.XML,
    )

    return agent


def run_workflow(prompt: str, persona_instructions: str) -> str:
    """
    Main function to set up and run the workflow, handling any exceptions that may occur.

    Args:
        prompt (str): Question/task for the agent to accomplish.
        persona_instructions (str): Extra role text injected via Jinja placeholder.

    Returns:
        str: The content generated by the agent.
    """
    try:
        agent = create_agent()

        result = agent.run(
            input_data={
                "input": prompt,
                "my_role": persona_instructions,
            },
        )

        return result.output.get("content")
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        return ""


if __name__ == "__main__":

    output = run_workflow(
        prompt=PROMPT,
        persona_instructions="Respond concisely and include the computed value in bold.",
    )

    logger.info("---------------------------------Result-------------------------------------")
    logger.info(output)
